# -*- coding: utf-8 -*-
"""
三国演义字符级RNN语言模型训练脚本
功能：从零开始训练一个RNN模型，学习《三国演义》文本风格，并能够生成类似文本
原理：使用循环神经网络(RNN)学习文本序列的统计规律，通过预测下一个字符的任务进行训练
"""
import os

import torch
import torch.nn as nn
import torch.optim as optim
import random

# https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_531.14_windows.exe
# pip install torch --index-url https://download.pytorch.org/whl/cu121 -v
# pip install torchaudio --index-url https://download.pytorch.org/whl/cu121 -v
# pip install torchvision --index-url https://download.pytorch.org/whl/cu121 -v
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("使用设备：", device)
print(torch.__version__)  # 查看torch版本
# 【1. 数据准备与预处理】
# 从文件中读取《三国演义》全文
with open("../../static/sanguo.txt", "r", encoding="utf-8") as f:
    text = f.read()  # 读取整个文本内容到字符串变量中

# 创建词汇表（所有唯一字符的集合）
chars = sorted(list(set(text)))  # 获取文本中所有不重复字符并排序
vocab_size = len(chars)  # 计算词汇表大小（有多少个不同的字符）
print(f"文本总长度: {len(text)} 字符, 唯一字符数: {vocab_size}")

# 创建字符到索引的映射字典（将每个字符映射为一个数字）
char_to_idx = {ch: i for i, ch in enumerate(chars)}  # 字符 -> 索引
# 创建索引到字符的映射字典（将数字映射回对应的字符）
idx_to_char = {i: ch for i, ch in enumerate(chars)}  # 索引 -> 字符

# 将整个文本转换为数字序列（模型只能处理数字）
data = [char_to_idx[ch] for ch in text]  # 将每个字符转换为其对应的索引值

# 【2. 模型超参数设置】
hidden_size = 128  # RNN隐藏状态的维度（记忆容量大小，越大模型能力越强但计算量越大）
seq_length = 100  # 训练时每个序列的长度（RNN展开的时间步数）
batch_size = 64  # 每次训练使用的样本数量（批量大小）
learning_rate = 0.002  # 学习率（控制参数更新步长的大小）
num_epochs = 50000  # 训练轮数（整个数据集上迭代的次数）
print_interval = 500  # 每隔多少轮打印一次训练状态


# 【3. 定义RNN模型结构】
class ThreeKingdomsRNN(nn.Module):
    """三国演义RNN语言模型"""

    def __init__(self, vocab_size, hidden_size, num_layers=2):
        """
        初始化RNN模型
        参数:
            vocab_size: 词汇表大小
            hidden_size: 隐藏层维度
            num_layers: RNN层数
        """
        super(ThreeKingdomsRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        # 嵌入层：将字符索引转换为密集向量表示
        # 作用：将离散的字符索引转换为连续的向量空间中的点，让模型能够学习字符的语义关系
        self.embedding = nn.Embedding(vocab_size, hidden_size)

        # RNN层：处理序列数据的核心层
        # 输入维度: hidden_size (嵌入后的向量维度)
        # 输出维度: hidden_size (隐藏状态的维度)
        # batch_first=True: 输入数据的维度为 (batch_size, seq_length, input_size)
        self.rnn = nn.RNN(
            hidden_size, hidden_size, num_layers, batch_first=True, nonlinearity="tanh"
        )

        # 输出层：将RNN的输出映射回词汇表空间
        # 作用：将隐藏状态转换为对下一个字符的预测概率分布
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden):
        """
        前向传播过程
        参数:
            x: 输入序列 (batch_size, seq_length)
            hidden: 初始隐藏状态
        返回:
            output: 每个时间步的输出 (batch_size, seq_length, vocab_size)
            hidden: 最终的隐藏状态，传递给下一个序列
        """
        # 1. 通过嵌入层将字符索引转换为向量
        out = self.embedding(x)  # 形状: (batch_size, seq_length, hidden_size)

        # 2. 通过RNN层处理序列
        # RNN返回:
        # - out: 每个时间步的隐藏状态 (batch_size, seq_length, hidden_size)
        # - hidden: 最后一个时间步的隐藏状态 (num_layers, batch_size, hidden_size)
        out, hidden = self.rnn(out, hidden)

        # 3. 通过全连接层将RNN输出映射回词汇表空间
        # 我们需要将输出重塑为 (batch_size * seq_length, hidden_size)
        # 然后通过全连接层得到 (batch_size * seq_length, vocab_size)
        batch_size, seq_length = x.size(0), x.size(1)
        out = out.contiguous().view(-1, self.hidden_size)
        out = self.fc(out)
        out = out.view(batch_size, seq_length, -1)

        return out, hidden

    def init_hidden(self, batch_size):
        """
        初始化隐藏状态（记忆）
        参数:
            batch_size: 批量大小
        返回:
            初始化为零的隐藏状态张量
        """
        # 创建一个全零的张量作为初始隐藏状态
        # 形状: (num_layers, batch_size, hidden_size)
        return torch.zeros(self.num_layers, batch_size, self.hidden_size)


def begin():
    # 创建模型实例
    model = ThreeKingdomsRNN(vocab_size, hidden_size).to(device)
    print(f"模型参数总量: {sum(p.numel() for p in model.parameters())}")

    # 【4. 定义损失函数和优化器】
    # 交叉熵损失函数：用于衡量预测概率分布与真实分布之间的差异
    criterion = nn.CrossEntropyLoss()

    # Adam优化器：用于更新模型参数，基于计算出的梯度
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 【5. 训练循环】
    print("开始训练RNN语言模型...")
    start_epoch = 0
    checkpoint_path = "../../static/sanguo_rnn.pth"
    if os.path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location=device)
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        start_epoch = checkpoint["epoch"]
        print(f"已加载断点: {checkpoint_path}，从第{start_epoch}轮继续训练")
    else:
        print("未检测到断点文件，从头开始训练")
    for epoch in range(start_epoch, num_epochs):
        # 5.1 准备训练数据
        # 随机选择批量训练的起始位置
        start_idxs = [
            random.randint(0, len(data) - seq_length - 1) for _ in range(batch_size)
        ]

        # 创建输入序列和目标序列
        # 输入序列: 从start_idx到start_idx+seq_length-1
        # 目标序列: 从start_idx+1到start_idx+seq_length (比输入序列右移一位)
        inputs = torch.zeros(batch_size, seq_length, dtype=torch.long).to(device)
        targets = torch.zeros(batch_size, seq_length, dtype=torch.long).to(device)

        for i, start_idx in enumerate(start_idxs):
            inputs[i] = torch.tensor(data[start_idx : start_idx + seq_length])
            targets[i] = torch.tensor(data[start_idx + 1 : start_idx + seq_length + 1])

        # 5.2 初始化隐藏状态
        hidden = model.init_hidden(batch_size).to(device)

        # 5.3 前向传播
        # 将输入传递给模型，得到输出和新的隐藏状态
        outputs, hidden = model(inputs, hidden)

        # 5.4 计算损失
        # 需要将输出和目标重塑为二维张量，以便计算损失
        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))

        # 5.5 反向传播
        # 清零梯度 -> 计算梯度 -> 更新参数
        optimizer.zero_grad()  # 清零梯度，防止梯度累积
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 更新模型参数

        # 5.6 定期打印训练状态
        if (epoch + 1) % print_interval == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

            # 生成一段文本示例，展示当前模型能力
            with torch.no_grad():  # 禁用梯度计算，节省内存
                # 使用随机字符作为起始
                start_char = random.choice(chars)
                input_seq = torch.tensor([[char_to_idx[start_char]]]).to(device)
                hidden = model.init_hidden(1).to(device)

                # 生成文本
                generated_text = start_char
                for _ in range(49):  # 生成50个字符
                    output, hidden = model(input_seq, hidden)
                    # 获取概率分布并采样下一个字符
                    prob = output.squeeze().exp()
                    char_idx = torch.multinomial(prob, 1).item()
                    # 将预测的字符添加到生成文本中
                    generated_text += idx_to_char[char_idx]
                    # 将预测的字符作为下一个输入
                    input_seq = torch.tensor([[char_idx]]).to(device)

                print(f"生成示例: {generated_text}")
            checkpoint = {
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "loss": loss.item(),
            }
            torch.save(checkpoint, f"../../static/sanguo_rnn.pth")
            print(f"已保存断点: sanguo_rnn.pth")


# 【6. 文本生成函数】
def generate_text(model, start_str, length=100, temperature=0.8):
    """
    使用训练好的模型生成文本
    参数:
        model: 训练好的模型
        start_str: 起始字符串
        length: 要生成的文本长度
        temperature: 控制生成随机性的参数（越高越随机，越低越保守）
    """
    model = model.to(device)
    model.eval()  # 设置模型为评估模式
    with torch.no_grad():  # 禁用梯度计算
        # 初始化输入和隐藏状态
        input_seq = torch.tensor([[char_to_idx[ch] for ch in start_str]]).to(device)
        hidden = model.init_hidden(1).to(device)

        # 生成文本
        generated_text = start_str
        for _ in range(length):
            output, hidden = model(input_seq, hidden)
            # 获取最后一个时间步的输出
            last_output = output[0, -1, :]
            # 应用温度参数调整概率分布
            scaled_logits = last_output / temperature
            # 将logits转换为概率分布
            probs = torch.softmax(scaled_logits, dim=0)
            # 从概率分布中采样下一个字符
            char_idx = torch.multinomial(probs, 1).item()
            # 将预测的字符添加到生成文本中
            generated_text += idx_to_char[char_idx]
            # 更新输入序列为最新预测的字符
            input_seq = torch.tensor([[char_idx]]).to(device)

        return generated_text


def test():
    checkpoint = torch.load("../../static/sanguo_rnn.pth", map_location=device)
    model = ThreeKingdomsRNN(vocab_size, hidden_size).to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    # model.load_state_dict(torch.load("../../static/sanguo_rnn.pth"))
    model.eval()
    # 使用训练好的模型生成文本
    print("\n使用训练好的模型生成文本:")
    print(generate_text(model, "话说天下大势", length=200))
    # print(generate_text(model, "曹操", length=100))


test()
# begin()
# print("训练完成!")
